# 导入必要的库
import pytorch_lightning as pl
import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import transforms, models


# 数据模块：负责加载和预处理数据
class TrafficDataModule(pl.LightningDataModule):
    def __init__(self, data_dir, batch_size=32):
        super().__init__()
        self.data_dir = data_dir       # 数据集路径
        self.batch_size = batch_size   # 批处理大小

        # 训练集数据增强和标准化
        self.train_transform = transforms.Compose([
            transforms.Resize(256),          # 调整尺寸
            transforms.RandomCrop(224),      # 随机裁剪
            transforms.RandomHorizontalFlip(), # 随机水平翻转
            transforms.ToTensor(),           # 转为张量
            transforms.Normalize(             # 标准化
                [0.485, 0.456, 0.406],
                [0.229, 0.224, 0.225])
        ])

        # 验证集预处理（无数据增强）
        self.val_transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),       # 中心裁剪
            transforms.ToTensor(),
            transforms.Normalize(
                [0.485, 0.456, 0.406],
                [0.229, 0.224, 0.225])
        ])

    # 数据准备阶段
    def setup(self, stage=None):
        # 使用ImageFolder自动加载分类数据集
        dataset = ImageFolder(self.data_dir, transform=self.train_transform)

        # 划分训练集和验证集（8:2比例）
        train_size = int(0.8 * len(dataset))
        self.train_ds, self.val_ds = torch.utils.data.random_split(
            dataset, [train_size, len(dataset)-train_size])

        # 为验证集应用不同的预处理
        self.val_ds.dataset.transform = self.val_transform

    # 训练数据加载器
    def train_dataloader(self):
        return DataLoader(
            self.train_ds,
            batch_size=self.batch_size,
            shuffle=True,    # 训练时打乱数据
            num_workers=4     # 使用4个进程加载数据
        )

    # 验证数据加载器
    def val_dataloader(self):
        return DataLoader(
            self.val_ds,
            batch_size=self.batch_size,
            num_workers=4
        )


# 定义分类模型
class ResNetClassifier(pl.LightningModule):
    def __init__(self, num_classes=10):
        super().__init__()
        # 加载预训练的ResNet-18
        # 加载预训练的ResNet-18
        self.model = models.resnet18(pretrained=False)  # 关闭自动下载
        # 加载本地预训练权重
        self.model.load_state_dict(torch.load('resnet18.pth'))
        # 替换最后一层全连接层
        # 替换最后一层全连接层
        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)
        # 定义交叉熵损失
        self.criterion = nn.CrossEntropyLoss()

    # 前向传播
    def forward(self, x):
        return self.model(x)

    # 训练步骤
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)                # 前向传播
        loss = self.criterion(y_hat, y) # 计算损失
        self.log('train_loss', loss)    # 记录训练损失
        return loss

    # 验证步骤
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)
        # 计算准确率
        acc = (y_hat.argmax(dim=1) == y).float().mean()
        # 记录验证指标（显示在进度条）
        self.log('val_acc', acc, prog_bar=True)
        self.log('val_loss', loss, prog_bar=True)

    # 配置优化器
    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=0.001)  # 使用Adam优化器


# 主程序
if __name__ == "__main__":
    # 初始化数据模块（使用实际数据集路径）
    dm = TrafficDataModule(r"F:\人工智能教材编写\traffic_sign\train")

    # 创建模型实例（类别数设为11）
    model = ResNetClassifier(num_classes=11)

    # 配置训练器
    trainer = pl.Trainer(
        max_epochs=2,  # 训练2个epoch
        callbacks=[
            # 模型检查点（保存最佳模型）
            pl.callbacks.ModelCheckpoint(
                monitor='val_acc',  # 监控验证准确率
                mode='max'          # 保存最大值
            )
        ],
        accelerator='auto'  # 自动选择GPU/CPU
    )

    # 开始训练
    trainer.fit(model, dm)

    # 保存最终模型权重
    torch.save(model.model, 'resnet18_traffic.pth')
